from ais.core import *
import time
import sys
import os
import copy
import logging
from pathlib import Path
from enum import Enum,unique
import torch
import numpy as np

from ais import __version__ as ais_version
from ais.utils import cfg, load_cfg_file, set_random_seed, get_state_dict, setup_logger, print_torch_version, logger_addHandler
from app.cobb_infer import DRCobbDet
from ais.image import  cv_imread, is_image_file, is_dcm_file, read_dcm_to_8bit, cv_show, cv_keyboard_run, cv_resize_shape, cv_imwrite

logger = None


@unique
class StucturalType(Enum):
    Nonstuctural = 0
    Stuctural = 1


class AisApp(object):
    def __init__(self, logger_file):
        """

        :param cfg_file: cfg file path
        :param logger_file: logger file path
        """
        logger = setup_logger(Path(__file__).name, level=logging.INFO, filepath=logger_file)
        self.DRCobbApp = DRCobbDet(logger_file)

    def init(self):
        """
        初始化模型以及相关工具
        :return:
        """
        self.DRCobbApp.init()

    def lenke_det(self, image_dict:dict):
        """
        DR lenke分型检测
        :param image_dict:
        {
        ap: "NPImage or None"
        lat: "NPImage or None"
        lbend: "NPImage or None"
        rbend: "NPImage or None"
        }
        :return: aisInfo
        """
        ais_info = []
        if 'ap' in image_dict:
            if image_dict.get('lbend', None) or image_dict.get('rbend', None) or image_dict.get('lat', None):
                ap_cobb_dict, vertebra_kp = self.DRCobbApp.process(image_dict['ap'], orientation='AP')
                orientation_dict = {"orientation": 'AP'}
                ais_info.append({**orientation_dict, **ap_cobb_dict})
                if image_dict['lbend'] or image_dict['rbend'] or image_dict['lat']:
                    # 对于主弯Cobb角小于30度的轻度侧弯,不进行三个弯的划分
                    # max_ap_angle = 0
                    # for k, v in ap_cobb_dict.items():
                    #     if v['angle'] > max_ap_angle:
                    #         max_ap_angle = v['angle']
                    # if max_ap_angle >= 25:
                    lbend_cobb_dict = None
                    rbend_cobb_dict = None
                    lat_cobb_dict = None
                    add_para = {'ap_dict': ap_cobb_dict}
                    if 'lbend' in image_dict:
                        lbend_cobb_dict = self.DRCobbApp.process(image_dict['lbend'], orientation='LBend', **add_para)
                        orientation_dict = {"orientation": 'LBend'}
                        ais_info.append({**orientation_dict, **lbend_cobb_dict})
                    if 'rbend' in image_dict:
                        rbend_cobb_dict = self.DRCobbApp.process(image_dict['rbend'], orientation='RBend', **add_para)
                        orientation_dict = {"orientation": 'RBend'}
                        ais_info.append({**orientation_dict, **rbend_cobb_dict})
                    if 'lat' in image_dict:
                        lat_cobb_dict = self.DRCobbApp.process(image_dict['lat'], orientation='LAT', lat_kp=vertebra_kp)
                        orientation_dict = {"orientation": 'LAT'}
                        ais_info.append({**orientation_dict, **lat_cobb_dict})
                    cobb_dict = {'ap': ap_cobb_dict, 'lbend': lbend_cobb_dict, 'rbend': rbend_cobb_dict, 'lat': lat_cobb_dict}
                    curve_type = self.ais_curve_bend(cobb_dict)

                    # PT、MT、TL/L 的 柔韧性计算
                    if 'lbend' in image_dict and 'rbend' in image_dict:
                        # PT flexibility
                        # ap_pt = ais_info[0]['PT']['angle']
                        # lbend_pt = ais_info[1]['PT']['angle']
                        # rbend_pt = ais_info[2]['PT']['angle']
                        if abs(ais_info[0]['PT']['angle']) > ais_info[1]['PT']['angle']:
                            pt_flex = round((abs(ais_info[0]['PT']['angle']) - ais_info[1]['PT']['angle']) / abs(ais_info[0]['PT']['angle']), 2)
                        elif abs(ais_info[0]['PT']['angle']) > ais_info[2]['PT']['angle']:
                            pt_flex = round((abs(ais_info[0]['PT']['angle']) - ais_info[2]['PT']['angle']) / abs(ais_info[0]['PT']['angle']), 2)
                        # MT flexibility
                        # ap_mt = ais_info[0]['MT']['angle']
                        # lbend_mt = ais_info[1]['MT']['angle']
                        # rbend_mt = ais_info[2]['MT']['angle']
                        if abs(ais_info[0]['MT']['angle']) > ais_info[1]['MT']['angle']:
                            mt_flex = round((abs(ais_info[0]['MT']['angle']) - ais_info[1]['MT']['angle']) / abs(ais_info[0]['MT']['angle']), 2)
                        elif abs(ais_info[0]['MT']['angle']) > ais_info[2]['MT']['angle']:
                            mt_flex = round((abs(ais_info[0]['MT']['angle']) - ais_info[2]['MT']['angle']) / abs(ais_info[0]['MT']['angle']), 2)
                        # TL/L flexibility
                        if abs(ais_info[0]['TL/L']['angle']) > ais_info[1]['TL/L']['angle']:
                            tl_flex = round((abs(ais_info[0]['TL/L']['angle']) - ais_info[1]['TL/L']['angle']) / abs(ais_info[0]['TL/L']['angle']), 2)
                        elif abs(ais_info[0]['TL/L']['angle']) > ais_info[2]['TL/L']['angle']:
                            tl_flex = round((abs(ais_info[0]['TL/L']['angle']) - ais_info[2]['TL/L']['angle']) / abs(ais_info[0]['TL/L']['angle']), 2)
                        ais_info.append({"flexibility": {"PT": pt_flex, 'MT': mt_flex, 'TL/L': tl_flex}})
                    ais_info.append({"orientation": "LENKE", "curve_type": curve_type, 'ls_modifier': '', 'ts_modifier': ''})
            else:
                ap_cobb_dict, _ = self.DRCobbApp.process(image_dict['ap'], orientation='AP')
                # curve_type = self.ais_curve_ap(ap_cobb_dict)
                orientation_dict = {"orientation": 'AP'}
                ais_info.append({**orientation_dict, **ap_cobb_dict})

        return ais_info

    def lenke_curve(self, curve_stuctural_dict: dict):
        """
        判断lenke侧凸
        cite: https://mp.weixin.qq.com/s/NsnnLWgQeDlkbfdBPNOnYA
        :param curve_stuctural_dict:
        :return:
        """
        if curve_stuctural_dict['TL/L'] == StucturalType.Nonstuctural:
            # type1 type2
            if curve_stuctural_dict['PT'] == StucturalType.Nonstuctural:
                return 1
            elif curve_stuctural_dict['PT'] == StucturalType.Stuctural:
                return 2
        else:
            # type3 type4 type5 type 6
            if curve_stuctural_dict['TL/L'] == StucturalType.Stuctural and curve_stuctural_dict['major']=='TL/L':
                # type5 type 6
                if curve_stuctural_dict['MT'] == StucturalType.Nonstuctural:
                    return 5
                elif curve_stuctural_dict['MT'] == StucturalType.Stuctural:
                    return 6
            else:
                # type3 type4
                # 'mt': Stuctural major
                if curve_stuctural_dict['PT'] == StucturalType.Nonstuctural:
                    return 3
                elif curve_stuctural_dict['PT'] == StucturalType.Stuctural:
                    return 4

    def ais_curve_bend(self, cobb_dict):
        """
        计算脊柱侧凸分型
        :param cobb_dict: {'ap': ap_cobb_dict, 'lbend': lbend_cobb_dict, 'rbend': rbend_cobb_dict}
        :return:
        """
        assert 'ap' in cobb_dict
        thoracic_index: list = ['PT', 'MT', 'TL/L']
        ap_angels:list = [cobb_dict['ap']['PT']['angle'], cobb_dict['ap']['MT']['angle'], cobb_dict['ap']['TL/L']['angle']]
        max_thoracic_index = thoracic_index[ap_angels.index(max(ap_angels))]

        curve_stuctural_dict ={}
        # curve_stuctural_dict = {max_thoracic_index: StucturalType.Stuctural}
        curve_stuctural_dict['major'] = max_thoracic_index

        # old
        # if 'rbend' in cobb_dict and 'lbend' in cobb_dict and cobb_dict['rbend'] and cobb_dict['lbend']:
        #     rbend_angels = [cobb_dict['rbend']['PT']['angle'], cobb_dict['rbend']['MT']['angle'], cobb_dict['rbend']['TL/L']['angle']]
        #     lbend_angels = [cobb_dict['lbend']['PT']['angle'], cobb_dict['lbend']['MT']['angle'], cobb_dict['lbend']['TL/L']['angle']]
        #     min_angles = np.min(np.vstack((rbend_angels, lbend_angels)), axis=0).tolist()
        #     for i, angle in enumerate(min_angles):
        #         if i != thoracic_index.index(max_thoracic_index):
        #             if angle >= 25.0:
        #                 curve_stuctural_dict[thoracic_index[i]] = StucturalType.Stuctural
        #             else:
        #                 curve_stuctural_dict[thoracic_index[i]] = StucturalType.Nonstuctural
        # elif 'rbend' in cobb_dict or 'lbend' in cobb_dict:
        #     if 'rbend' in cobb_dict and cobb_dict['rbend']:
        #         bend_angels = [cobb_dict['rbend']['PT']['angle'], cobb_dict['rbend']['MT']['angle'], cobb_dict['rbend']['TL/L']['angle']]
        #         for i, angle in enumerate(bend_angels):
        #             if i != thoracic_index.index(max_thoracic_index):
        #                 if angle >= 25.0:
        #                     curve_stuctural_dict[thoracic_index[i]] = StucturalType.Stuctural
        #                 else:
        #                     curve_stuctural_dict[thoracic_index[i]] = StucturalType.Nonstuctural
        #
        #     elif 'lbend' in cobb_dict and cobb_dict['lbend']:
        #         bend_angels = [cobb_dict['lbend']['PT']['angle'], cobb_dict['lbend']['MT']['angle'], cobb_dict['lbend']['TL/L']['angle']]
        #         for i, angle in enumerate(bend_angels):
        #             if i != thoracic_index.index(max_thoracic_index):
        #                 if angle >= 25.0:
        #                     curve_stuctural_dict[thoracic_index[i]] = StucturalType.Stuctural
        #                 else:
        #                     curve_stuctural_dict[thoracic_index[i]] = StucturalType.Nonstuctural

        if 'rbend' in cobb_dict and 'lbend' in cobb_dict and cobb_dict['rbend'] and cobb_dict['lbend']:
            rbend_angels = [cobb_dict['rbend']['PT']['angle'], cobb_dict['rbend']['MT']['angle'], cobb_dict['rbend']['TL/L']['angle']]
            lbend_angels = [cobb_dict['lbend']['PT']['angle'], cobb_dict['lbend']['MT']['angle'], cobb_dict['lbend']['TL/L']['angle']]
            min_angles = np.min(np.vstack((rbend_angels, lbend_angels)), axis=0).tolist()
            for i, angle in enumerate(min_angles):
                if angle >= 25.0:
                    curve_stuctural_dict[thoracic_index[i]] = StucturalType.Stuctural
                else:
                    curve_stuctural_dict[thoracic_index[i]] = StucturalType.Nonstuctural

            # 如果三个弯都不是结构性弯，那主弯就认为是结构性弯
            if curve_stuctural_dict['PT'] == StucturalType.Nonstuctural and curve_stuctural_dict['MT'] == StucturalType.Nonstuctural  \
                    and curve_stuctural_dict['TL/L'] == StucturalType.Nonstuctural:
                curve_stuctural_dict[max_thoracic_index] = StucturalType.Stuctural

        elif 'rbend' in cobb_dict or 'lbend' in cobb_dict:
            if 'rbend' in cobb_dict and cobb_dict['rbend']:
                bend_angels = [cobb_dict['rbend']['PT']['angle'], cobb_dict['rbend']['MT']['angle'], cobb_dict['rbend']['TL/L']['angle']]
                for i, angle in enumerate(bend_angels):
                    if angle >= 25.0:
                        curve_stuctural_dict[thoracic_index[i]] = StucturalType.Stuctural
                    else:
                        curve_stuctural_dict[thoracic_index[i]] = StucturalType.Nonstuctural

            elif 'lbend' in cobb_dict and cobb_dict['lbend']:
                bend_angels = [cobb_dict['lbend']['PT']['angle'], cobb_dict['lbend']['MT']['angle'], cobb_dict['lbend']['TL/L']['angle']]
                for i, angle in enumerate(bend_angels):
                    if angle >= 25.0:
                        curve_stuctural_dict[thoracic_index[i]] = StucturalType.Stuctural
                    else:
                        curve_stuctural_dict[thoracic_index[i]] = StucturalType.Nonstuctural
        return self.lenke_curve(curve_stuctural_dict)

    def read_image(self, image_or_path):
        """
        读取普通图像或者DICOM图像
        :param image_or_path:
        :return:
        """
        if isinstance(image_or_path, NPArray):
            show_pred_image = image_or_path
        elif isinstance(image_or_path, Path) or isinstance(str(image_or_path), str):
            if not Path(image_or_path).exists():
                logger.info(__name__ + " process input: {} file not exist!".format(str(image_or_path)))
                return None
            try:
                if is_dcm_file(str(image_or_path)):
                    show_pred_image = read_dcm_to_8bit(str(image_or_path))
                    show_pred_image = np.repeat(show_pred_image[..., np.newaxis], 3, 2)
                elif is_image_file(str(image_or_path)):
                    show_pred_image = cv_imread(str(image_or_path), 1)
            except Exception as e:
                logger.error(__name__ + " process input: {} file not exist!".format(str(image_or_path)))
                return None
        return show_pred_image

    def evaluate_case(self, case_dir:Path):
        """AIS Lenke测量本地测试评估"""
        import cv2
        # ap_file = sorted(case_dir.joinpath('ap').glob('*.jpg'))[0]
        # lat_file = sorted(case_dir.joinpath('lat').glob('*.jpg'))[0]
        # lbend_file = sorted(case_dir.joinpath('lbend').glob('*.jpg'))[0]
        # rbend_file = sorted(case_dir.joinpath('rbend').glob('*.jpg'))[0]
        ap_file = case_dir.joinpath('ap.jpg')
        lat_file = case_dir.joinpath('lat.jpg')
        lbend_file = case_dir.joinpath('lbend.jpg')
        rbend_file = case_dir.joinpath('rbend.jpg')

        image_dict = {'ap': ap_file,
                      'lat': lat_file,
                      'lbend': lbend_file,
                      'rbend': rbend_file}
        ais_info = self.lenke_det(image_dict)

        # show
        ap_image = self.read_image(ap_file)
        lat_image = self.read_image(lat_file)
        lbend_image = self.read_image(lbend_file)
        rbend_image = self.read_image(rbend_file)

        def draw_info(image, ais_dict, key):
            ais_dict_copy = copy.deepcopy(ais_dict)
            for cobb_dict in ais_dict_copy:
                if 'orientation' in cobb_dict:
                    if cobb_dict['orientation'] == key:
                        del cobb_dict['orientation']
                        for k, v in cobb_dict.items():
                            cv2.circle(image, (v['up_line']['start_x'], v['up_line']['start_y']), 1, (0, 0, 255), 4)
                            cv2.circle(image, (v['up_line']['end_x'], v['up_line']['end_y']), 1, (0, 0, 255), 4)
                            cv2.line(image, (v['up_line']['start_x'], v['up_line']['start_y']),
                                     (v['up_line']['end_x'], v['up_line']['end_y']), (0, 0, 255), 3)

                            cv2.circle(image, (v['down_line']['start_x'], v['down_line']['start_y']), 1, (0, 255, 0), 4)
                            cv2.circle(image, (v['down_line']['end_x'], v['down_line']['end_y']), 1, (0, 255, 0), 4)
                            cv2.line(image, (v['down_line']['start_x'], v['down_line']['start_y']),
                                     (v['down_line']['end_x'], v['down_line']['end_y']), (0, 255, 0), 3)
                        show_cobb = ""
                        if 'PT' in cobb_dict:
                            show_cobb += 'PT (' + cobb_dict['PT']['up_endplate'] + '-' + cobb_dict['PT'][
                                'down_endplate'] + ') : ' + str(cobb_dict['PT']['angle']) + '\n'
                        if 'MT' in cobb_dict:
                            show_cobb += 'MT (' + cobb_dict['MT']['up_endplate'] + '-' + cobb_dict['MT'][
                                'down_endplate'] + ') : ' + str(cobb_dict['MT']['angle']) + '\n'
                        if 'TL/L' in cobb_dict:
                            show_cobb += 'TL/L (' + cobb_dict['TL/L']['up_endplate'] + '-' + cobb_dict['TL/L'][
                                'down_endplate'] + ') : ' + str(cobb_dict['TL/L']['angle']) + '\n'
                        text_cobb = show_cobb.split("\n")
                        for i, text in enumerate(text_cobb[::-1]):
                            if len(text) > 1:
                                image = cv2.putText(image, text, (15, ap_image.shape[0] - i * 40 - 80),
                                                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                if key == 'AP':
                    if 'LENKE' in cobb_dict:
                        if 'curve_type' in cobb_dict['LENKE'] and cobb_dict['LENKE']['curve_type']:
                            text = "Det AIS Curve Type: " + str(cobb_dict['LENKE']['curve_type'])
                            image = cv2.putText(image, text, (15, ap_image.shape[0] - 4 * 40 - 80),
                                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

                            gt_curve = case_dir.parent.name
                            text = "Gt AIS Curve Type: " + str(gt_curve)
                            image = cv2.putText(image, text, (15, ap_image.shape[0] - 5 * 40 - 80),
                                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                        return image

            return image

        # draw ap cobb image
        ap_image = draw_info(ap_image, ais_info, 'AP')
        lat_image = draw_info(lat_image, ais_info, 'LAT')
        lbend_image = draw_info(lbend_image, ais_info, 'LBend')
        rbend_image = draw_info(rbend_image, ais_info, 'RBend')
        lbend_image = cv_resize_shape(lbend_image, ap_image.shape[0:2])
        rbend_image = cv_resize_shape(rbend_image, ap_image.shape[0:2])
        if lat_image is not None:
            lat_image = cv_resize_shape(lat_image, ap_image.shape[0:2])

        cv_show('ais', np.hstack((ap_image, lat_image, lbend_image, rbend_image)), resize=0.9)
        save_name = case_dir.joinpath(case_dir.name +'_ais.jpg')
        cv_imwrite(str(save_name), np.hstack((ap_image, lat_image, lbend_image, rbend_image)))
        cv_keyboard_run()


def test_AisApp():
    import argparse, configparser

    def parse_args():
        parser = argparse.ArgumentParser(description="AIS DR Measurement Inference Service!")
        parser.add_argument('--ini',
                            type=str,
                            help='ini config file',
                            default='web_server.ini')
        return parser.parse_args()

    def parse_config_ini(ini_file: Path):
        if Path(ini_file).exists():
            config = configparser.ConfigParser()
            config.read(str(ini_file))
            global server_name, debug_flag, docker_data_root, host_data_root, deploy_config_file, image_url_prefix, port_number
            server_name = config['DEFAULT']['server_name']
            debug_flag = config['DEFAULT'].getboolean('debug_flag')
            port_number = int(config['DEFAULT']['port_number'])
            docker_data_root = config['docker']['docker_data_root']
            host_data_root = config['docker']['host_data_root']
            deploy_config_file = config['ais_model']['deploy_config_file']
            image_url_prefix = config['backend_service']['image_url_prefix']
        else:
            print("ini_file: {} not found!!!".format(str(ini_file)))
            sys.exit(1)
    # parse ini file
    args = parse_args()
    parse_config_ini(args.ini)

    # load model deploy_config_file
    deploy_config_file = '/home/cao/PycharmProjects/AIS_Algorithm/experiment/app/ais_seg_hrnet_w18_landmark_spinnet_res34_v1.0.0.yml'
    load_cfg_file(deploy_config_file)
    logger_file = Path(cfg.LOG.DIR).joinpath("AIS-WebServer" + time.strftime("--%Y-%m-%d-%H-%M-%S", time.localtime()))
    if not Path(logger_file.parent).exists():
        os.makedirs(str(logger_file.parent))

    global logger, ais_app
    logger = setup_logger(Path(__file__).name, level=logging.INFO, filepath=logger_file)

    logger.info("start init ais_app...")
    ais_app = AisApp(logger_file)
    ais_app.init()
    logger.info("init ais_app successed!")

    data_root = '/home/cao/data/medical/datasets/vertebral/AIS/ais'
    data_path = Path(data_root)
    for i in range(1, 7):
        image_items = sorted(data_path.joinpath('data', 'all', 'lenke'+str(5)).glob('*'))

        for image_dir in image_items:
            print(image_dir)
            image_dir = image_dir.with_name('haugnjiao')
            if image_dir.is_dir():
                ais_app.evaluate_case(image_dir)


if __name__ == '__main__':
    print(__name__)
    test_AisApp()





